from mbkEGES.Random_walk import R_walk
from collections import defaultdict
from tqdm import tqdm
import random
import torch

class Getsample():
    def __init__(self, Grapher, walklength, walknumber, window = 3, negative_rate = 1, one_or_two = 0):
        """
        Graph: 提前加载的grapher类对象
        walklength: 随机游走的长度
        walknumber: 每个节点随机游走的次数
        window: sikpgram算法中的窗口大小
        negative_rate: 负样例/正样例的大小
        one_or_two: 默认0, 选择单向图还是双向图
        """
        self.g = Grapher
        self.G = [self.g.one_way_G(), self.g.two_way_G()][one_or_two]
        self.window = window
        self.neg_rate = negative_rate
        self.random_walker = R_walk(self.G, walklength, walknumber)
        self.sequences = self.random_walker.random_walk()
        self.all_node = set(self.G.keys())

    def positive_negative_sample(self):

        pos_sample = [] #存放正样例
        neg_sample = [] #存放负样例
        for sentence in tqdm(self.sequences):
            for i, centerword in enumerate(sentence):
                left = max(0,i-self.window)
                right = min(i+self.window+1, len(sentence))
                skipword = sentence[left:i]+sentence[i+1:right]
                for word in skipword:
                    pos_sample.append([centerword, word]+ list(self.g.item_side_dict[centerword].values()) +[1]) #[center, contex, side_informations, 1]
                #获得所有可能的负样本集合  
                all_nega_sample = list(self.all_node - set(skipword))
                for j in range(int(2*self.window*self.neg_rate)):
                    neg_sample.append([centerword, random.choice(all_nega_sample)]+\
                        list(self.g.item_side_dict[centerword].values())+[0]) #[center, contex, side_informations, 0]
        print(f'正负样本总共有{len(pos_sample)+len(neg_sample)}.')
        return pos_sample+neg_sample

def integration(samples):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    centers = []
    contexs = []
    side_dim = len(samples[0])-3
    side_information_dict = defaultdict(list)
    side_information_list = []
    labels = []
    for i in range(len(samples)):
        centers.append(samples[i][0])
        contexs.append(samples[i][1])
        labels.append(samples[i][-1])
        for j in range(side_dim):
            side_information_dict[j].append(samples[i][2+j])
    for j in range(side_dim):
        side_information_list.append(torch.tensor(side_information_dict[j]).to(device))
    return torch.tensor(centers).to(device), torch.tensor(contexs).to(device), side_information_list, torch.tensor(labels).to(device)